import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
data = pd.read_csv('D:\\code\\rlearn\\stock_399344_var_model_data_v2.csv')
data['date'] = data['date'].astype('datetime64[ns]')
symbols = set(data['symbol'])

lva = "close_var_l_ma_fib_m1"
fva = "close_var_f_ma_fib_m1"
va = "close_var_ma_fib_m1"
turn = 'turnover_ma_fib_m1'
### 数据处理
df = data
## close_var 增长
df['close_var_ma3_upd'] =  df['close_var_ma3']  - df.groupby(['symbol'])['close_var_ma3'].shift(3)
df['close_var_ma5_upd'] =  df['close_var_ma5']  - df.groupby(['symbol'])['close_var_ma5'].shift(5)
df['close_var_ma10_upd'] = df['close_var_ma10'] - df.groupby(['symbol'])['close_var_ma10'].shift(10)

data[lva + "upd"] = data[lva] - data.groupby(['symbol'])[lva].shift(5)
data[fva + "upd"] = data[fva] - data.groupby(['symbol'])[fva].shift(5)
data[va + "upd"] = data[va] - data.groupby(['symbol'])[va].shift(5)

df[turn + "upd"] = df[turn] - df.groupby(['symbol'])[turn].shift(5)

data = data.dropna(axis=0, how='any')


sns.set_style('ticks')
sns.set_context('paper')
i = 0
count = 0
count_map = {}
q = {}
for symbol in symbols:
    print("----------" + str(i) + " :" + symbol + "-----------")
    i = i + 1
    is_buy = []
    mdata = data[data['symbol'] == symbol]
    scat = {
        "day": [],
        "price": [],
        "action": []
    }
    # 当var_l 上升的时候 上传了 var_f 买入
    # 当var_l 和 var_f 都小于 var 的时候卖出
    for index, row in mdata.iterrows():
        # print(index)
        if not q:
            q = {
                "close": row["close"],
                "date": row["date"],
                lva: row[lva],
                fva: row[fva]
            }
            continue
        else:
            old_row = q
            # if not is_buy and row[lva] > old_row[lva] and row[lva] > row[fva] and old_row[fva] > old_row[lva]:
            if not is_buy and row[lva + "upd"] > 0.07 and row[va + "upd"] > 0\
                    and row[fva + "upd"] > 0 and row[turn + "upd"] > 0:
                print("buy[" + symbol + ']: price:' + str(row["close"]) + " date:" + str(row['date']))
                is_buy = [symbol, row['date'], row["close"]]
                scat["day"].append(row['date'])
                scat["price"].append(row['close'])
                scat["action"].append('buy')
                continue

            # if is_buy and row[va] > row[lva] and row[va] > row[fva]:
            if is_buy and row[lva + "upd"] < -0.06:
                rate = int(round(100 * (row["close"] / is_buy[2] - 1)))
                if not count_map.__contains__(rate):
                    count_map[rate] = 1
                else:
                    count_map[rate] = count_map[rate] + 1
                count = count + 1
                print("sell[" + symbol + ']: price:' + str(row["close"]) + " date:" + str(row['date']) + " rate:" + str(rate))
                is_buy.append(row["close"])
                scat["day"].append(row['date'])
                scat["price"].append(row['close'])
                scat["action"].append('sell')
                is_buy = []
            q = {
                "close": row["close"],
                "date": row["date"],
                lva: row[lva],
                fva: row[fva]
            }
    if is_buy:
        old_row = q
        rate = int(round(100 * (old_row["close"] / is_buy[2] - 1)))
        if not count_map.__contains__(rate):
            count_map[rate] = 1
        else:
            count_map[rate] = count_map[rate] + 1
        count = count + 1
        print("sell[" + symbol + ']: price:' + str(old_row["close"]) + " date:" + str(old_row['date']) + " rate:" + str(rate))
        is_buy.append(old_row["close"])
        scat["day"].append(old_row['date'])
        scat["price"].append(old_row['close'])
        scat["action"].append('sell')
        is_buy = []
    q = {}
    #   画图
    f = plt.figure()
    f.add_subplot(3, 1, 1)
    plt.title(symbol)
    sns.scatterplot(pd.Series(scat["day"]), pd.Series(scat["price"]), hue=pd.Series(scat["action"]), palette="Set2")
    sns.lineplot(data=mdata, x='date', y='close')
    f.add_subplot(3, 1, 2)
    plt.title("var_ma")
    sns.lineplot(data=[mdata[va], mdata[fva], mdata[lva], mdata['turnover_ma_fib_m1']])
    f.add_subplot(3, 1, 3)
    plt.title("var_upd")
    sns.lineplot(data=[mdata["close_var_ma5_upd"], mdata["close_var_ma10_upd"]])
    # plt.savefig("D:\\code\\rlearn\\picture1\\" + symbol + ".png")
    plt.show()
    xxx = 1



print(count_map)
print(count)
for i in range(200):
    if i == 0:
        print(i, count_map.get(i))
    if i != 0:
        print(i, count_map.get(i), -i, count_map.get(-i))

### todo 止损，分步止损，6，9，12，15，抛完
### todo xgboost 看下这三个参数
### todo 买入条件 var var_f var_l turnover上涨比例
### todo 前n天值的最大值max 下降min(x, %25max)， 就卖出，  前n天的最小值 上涨多少区间就买入？
###  指标不断上涨，则要加仓，指标不断下降，则要减仓
###  不买涨停的股票
###  如何对比各个股票